from django.core.management.base import BaseCommand
from django.db import transaction
from movies.models import Movie, Genre, Rating
from django.contrib.auth.models import User
import csv
import os
import random
from datetime import datetime

class Command(BaseCommand):
    help = '从MovieLens数据集导入电影'

    def add_arguments(self, parser):
        parser.add_argument('--path', type=str, help='MovieLens数据集文件夹路径')
        parser.add_argument('--sample', type=int, default=1000, help='导入电影的数量')
        parser.add_argument('--ratings', type=int, default=5000, help='导入评分的数量')

    def handle(self, *args, **options):
        path = options['path']
        sample_size = options['sample']
        ratings_size = options['ratings']

        if not path or not os.path.exists(path):
            self.stdout.write(self.style.ERROR(f'路径不存在: {path}'))
            return

        # 确保有测试用户
        test_users = self.ensure_test_users()

        # 导入电影数据
        self.import_movies(path, sample_size)
        
        # 导入评分数据
        self.import_ratings(path, ratings_size, test_users)
        
        self.stdout.write(self.style.SUCCESS('数据导入完成!'))

    def ensure_test_users(self):
        """确保系统中有测试用户"""
        test_users = []
        for i in range(1, 11):
            username = f'user{i}'
            user, created = User.objects.get_or_create(
                username=username,
                defaults={
                    'email': f'{username}@example.com',
                }
            )
            if created:
                user.set_password(username)
                user.save()
                self.stdout.write(self.style.SUCCESS(f'创建测试用户: {username}'))
            test_users.append(user)
        return test_users

    def import_movies(self, path, sample_size):
        """导入电影数据"""
        movies_file = os.path.join(path, 'movies.csv')
        if not os.path.exists(movies_file):
            self.stdout.write(self.style.ERROR(f'电影数据文件不存在: {movies_file}'))
            return

        # 获取所有现有的类型
        existing_genres = {genre.name: genre for genre in Genre.objects.all()}

        with open(movies_file, 'r', encoding='utf-8') as f:
            reader = csv.DictReader(f)
            all_movies = list(reader)
            
            # 如果指定了样本大小，随机选择电影
            if sample_size > 0 and sample_size < len(all_movies):
                selected_movies = random.sample(all_movies, sample_size)
            else:
                selected_movies = all_movies

            with transaction.atomic():
                for row in selected_movies:
                    movie_id = row['movieId']
                    title_with_year = row['title']
                    
                    # 从标题中提取年份
                    year = None
                    if '(' in title_with_year and ')' in title_with_year:
                        title = title_with_year[:title_with_year.rfind('(')].strip()
                        year_str = title_with_year[title_with_year.rfind('(')+1:title_with_year.rfind(')')]
                        try:
                            year = int(year_str)
                        except ValueError:
                            year = 2000  # 默认年份
                    else:
                        title = title_with_year
                        year = 2000  # 默认年份
                    
                    # 处理类型
                    genre_names = row['genres'].split('|')
                    
                    # 如果电影已存在，跳过
                    if Movie.objects.filter(title=title, year=year).exists():
                        continue
                    
                    # 创建电影
                    movie = Movie.objects.create(
                        title=title,
                        year=year,
                        description=f"这是电影 {title} 的简介。",
                        poster_url=f"https://via.placeholder.com/300x450?text={title.replace(' ', '+')}"
                    )
                    
                    # 添加类型
                    for genre_name in genre_names:
                        if genre_name == '(no genres listed)':
                            continue
                            
                        if genre_name not in existing_genres:
                            genre = Genre.objects.create(name=genre_name)
                            existing_genres[genre_name] = genre
                        else:
                            genre = existing_genres[genre_name]
                        
                        movie.genres.add(genre)
                    
                    self.stdout.write(f'导入电影: {title} ({year})')

    def import_ratings(self, path, sample_size, test_users):
        """导入评分数据"""
        ratings_file = os.path.join(path, 'ratings.csv')
        if not os.path.exists(ratings_file):
            self.stdout.write(self.style.ERROR(f'评分数据文件不存在: {ratings_file}'))
            return

        # 获取所有电影的ID映射
        movie_ids = {str(m.id): m for m in Movie.objects.all()}
        
        with open(ratings_file, 'r', encoding='utf-8') as f:
            reader = csv.DictReader(f)
            ratings_to_add = []
            
            for row in reader:
                # 获取对应的电影
                movie_id = row['movieId']
                movie = Movie.objects.filter(title__icontains=movie_id).first()
                
                if not movie:
                    continue
                
                # 随机选择一个测试用户
                user = random.choice(test_users)
                
                # 获取评分和时间戳
                rating_value = min(5, max(1, int(float(row['rating']))))
                timestamp = int(row['timestamp'])
                rated_at = datetime.fromtimestamp(timestamp)
                
                # 如果已经有这个用户对这部电影的评分，跳过
                if Rating.objects.filter(user=user, movie=movie).exists():
                    continue
                
                # 创建评分
                rating = Rating(
                    user=user,
                    movie=movie,
                    rating=rating_value,
                    comment=f"这是用户 {user.username} 对电影 {movie.title} 的评价。",
                    created_at=rated_at
                )
                ratings_to_add.append(rating)
                
                if len(ratings_to_add) >= sample_size:
                    break
            
            # 批量创建评分
            with transaction.atomic():
                Rating.objects.bulk_create(ratings_to_add)
                
            self.stdout.write(f'导入了 {len(ratings_to_add)} 条评分') 